# -*- coding: utf-8 -*-
# time: 2025/5/10 09:18
# file: tf01.py
# author: hanson
"""
小模型推理
Pythia-70m需要约1GB GPU内存

"""
import torch
from modelscope import pipeline

from transformers import AutoTokenizer, AutoModelForCausalLM

from base.tansformers_pip import generator

tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m")
model = AutoModelForCausalLM.from_pretrained("EleutherAI/pythia-70m")

# 将模型移动到GPU（如果有的话）
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# 准备输入文本
prompt = "你知道中山大学在哪里吗"

# 编码输入文本
inputs = tokenizer(prompt, return_tensors="pt").to(device)

# 让模型推理并生成文本
outputs = model.generate(
    **inputs,
    max_new_tokens=50,  # 生成的最大新token数量
    do_sample=True,     # 是否使用采样
    temperature=0.7,    # 控制随机性的温度参数
    top_k=50,           # top-k采样参数
    top_p=0.95,         # top-p（核）采样参数
)

# 解码生成的文本
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

print(generated_text)

from transformers import pipeline
# 如果你想要更简单的生成方式，也可以使用pipeline：
pipeline_text_generation = pipeline("text-generation", model="EleutherAI/pythia-70m")
result = generator("请用中文介绍自己", max_length=50, do_sample=True)
print(result[0]["generated_text"])